import argparse
import os
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
import cv2
import numpy as np
from pathlib import Path

def print_stats(image, name="Image"):
    """Calculates and prints statistics for a given image."""
    if image is None:
        print(f"Cannot calculate stats for {name}, as it is None.")
        return

    # For multi-channel images, calculate stats across all channels
    min_val = np.min(image)
    max_val = np.max(image)
    mean_val = np.mean(image)
    median_val = np.median(image)
    std_val = np.std(image)

    print("-" * 30)
    print(f"Statistics for: {name}")
    print("-" * 30)
    print(f"  Shape: {image.shape}")
    print(f"  Data Type: {image.dtype}")
    print(f"  Min value:  {min_val:.4f}")
    print(f"  Max value:  {max_val:.4f}")
    print(f"  Mean value: {mean_val:.4f}")
    print(f"  Median value: {median_val:.4f}")
    print(f"  Std Dev:    {std_val:.4f}")
    print("-" * 30 + "\n")


def custom_tonemap(color: np.ndarray) -> np.ndarray:
    """
    Applies a custom tone mapping algorithm that scales the image
    brightness based on a percentile.
    """
    # Tone mapping parameters.
    gamma = 1.0 / 2.2  # Standard gamma correction exponent.
    inv_gamma = 1.0 / gamma
    percentile = 90  # Desired percentile brightness in the unmodified image.
    brightness_nth_percentile_desired = 0.8  # Desired brightness after scaling.

    # Assume all pixels are valid for this general purpose script
    valid_mask = np.ones(color.shape[:2], dtype=bool)

    # Calculate brightness (luminance)
    brightness = (
        0.3 * color[:, :, 0] + 0.59 * color[:, :, 1] + 0.11 * color[:, :, 2]
    )
    brightness_valid = brightness[valid_mask]

    eps = 0.0001  # Avoid division by zero.
    brightness_nth_percentile_current = np.percentile(
        brightness_valid, percentile
    )

    if brightness_nth_percentile_current < eps:
        scale = 0.0
    else:
        scale = (
            np.power(brightness_nth_percentile_desired, inv_gamma)
            / brightness_nth_percentile_current
        )

    # Apply scaling and gamma correction
    color_tonemapped = np.power(np.maximum(scale * color, 0), gamma)
    
    # Clip to the valid [0.0, 1.0] range
    color_tonemapped = np.clip(color_tonemapped, 0.0, 1.0)

    return color_tonemapped


def main():
    """Main function to parse arguments and run the processing pipeline."""
    parser = argparse.ArgumentParser(
        description="Process HDR (.exr) images, apply tone mapping, and save the results."
    )
    parser.add_argument(
        "--exr_path", 
        type=str, 
        default='E:/data/3D/infinigen/1a1f2004/frames/Image/camera_0/Image_0_0_0002_0.exr',
        help="Full path to the input .exr file."
    )
    parser.add_argument(
        "--output_dir", 
        type=str, 
        default="D:/liuco/Downloads/output",
        help="Path to the directory to save output files."
    )
    args = parser.parse_args()

    # --- 1. Setup Paths and Directories ---
    exr_path = Path(args.exr_path)
    output_dir = Path(args.output_dir)

    if not exr_path.is_file():
        print(f"Error: Input file not found at {exr_path}")
        return

    # Create output directory if it doesn't exist
    output_dir.mkdir(parents=True, exist_ok=True)
    print(f"Output will be saved to: {output_dir}\n")

    # --- 2. Read EXR and Print Stats ---
    # Load the HDR image preserving its high dynamic range (32-bit float)
    hdr_image = cv2.imread(str(exr_path), cv2.IMREAD_UNCHANGED)
    if hdr_image is None:
        print(f"Error: Failed to read EXR file at {exr_path}")
        return
    print_stats(hdr_image, f"Original EXR: {exr_path.name}")
    
    depth_path = Path(str(exr_path).replace("Image", "Depth")).with_suffix(".png")
    depth_image = cv2.imread(str(depth_path), cv2.IMREAD_UNCHANGED)
    if depth_image is None:
        print(f"Error: Failed to read depth file at {depth_path}")
        return
    for i in range(3):
        print_stats(depth_image[:, :, i], f"Original depth: {depth_path.name}")
    print_stats(depth_image.mean(axis=-1), f"Original depth mean: {depth_path.name}")

    depth_path = Path(str(exr_path).replace("Image", "Depth")).with_suffix(".npy")
    depth_image = np.load(depth_path).astype(np.float32)
    depth_image[~np.isfinite(depth_image)] = 0.0  # Handle invalid depths
    if depth_image is None:
        print(f"Error: Failed to read depth file at {depth_path}")
        return
    print_stats(depth_image, f"Original depth: {depth_path.name}")

    # --- 3. Read Corresponding PNG and Print Stats ---
    png_path = exr_path.with_suffix(".png")
    if not png_path.is_file():
        print(f"Warning: Corresponding PNG file not found at {png_path}\n")
    else:
        png_image = cv2.imread(str(png_path))
        cv2.imwrite(str(png_output_path), png_image)
        # Normalize from [0, 255] to [0.0, 1.0] for comparison
        # png_image_normalized = png_image.astype(np.float32) / 255.0
        print_stats(png_image, f"Corresponding PNG: {png_path.name}")

    if hdr_image.ndim == 3:
        hdr_image = cv2.cvtColor(hdr_image, cv2.COLOR_BGR2RGB)
    hdr_image = hdr_image.astype(np.float32)
    # --- 4. Apply Custom Tone Mapping ---
    print("Applying Custom Tone Mapping...")
    custom_mapped = custom_tonemap(hdr_image)
    
    # Convert from float [0, 1] to uint8 [0, 255] for saving
    custom_mapped_8bit = (custom_mapped * 255).astype(np.uint8)
    custom_output_path = output_dir / "custom.png"
    cv2.imwrite(str(custom_output_path), custom_mapped_8bit)
    print(f"Saved custom tone mapped image to: {custom_output_path}\n")
    
    # --- 5. Apply Various OpenCV Tone Mappings ---
    print("Applying OpenCV Tone Mapping methods...")
    
    # Dictionary of OpenCV tonemappers to apply
    opencv_tonemappers = {
        "reinhard": cv2.createTonemapReinhard(gamma=2.2),
        "drago": cv2.createTonemapDrago(gamma=2.2, saturation=0.7),
        "mantiuk": cv2.createTonemapMantiuk(gamma=2.2, scale=0.7)
    }

    for name, tonemap_algorithm in opencv_tonemappers.items():
        print(f"  Processing with {name.capitalize()}...")
        # Use .copy() as process can sometimes modify the image in-place
        sdr_image = tonemap_algorithm.process(hdr_image.copy())
        
        # Convert float output to 8-bit for saving
        sdr_image_8bit = np.clip(sdr_image * 255, 0, 255).astype(np.uint8)
        
        # Save the result
        output_filename = output_dir / f"opencv_{name}.png"
        cv2.imwrite(str(output_filename), sdr_image_8bit)
        print(f"    Saved to: {output_filename}")
        
    print("\nProcessing complete. ✨")


if __name__ == "__main__":
    main()